#include <gtest/gtest.h>
#include "base_memory_pool.h"
#include "base_log.h"
#include <vector>
#include <thread>
#include <random>
#include <algorithm>

using namespace El::Base;

class MemoryPoolTest : public ::testing::Test {
protected:
    void SetUp() override
    {
        // 每个测试前重置内存池
        MemoryPool::GetInstance().Reset();
    }

    void TearDown() override
    {
        // 每个测试后重置内存池
        MemoryPool::GetInstance().Reset();
        // 释放所有实例
        MemoryPool::ReleaseAllInstances();
    }

    // 生成随机数据
    std::vector<uint8_t> GenerateRandomData(size_t size)
    {
        std::vector<uint8_t> data(size);
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> distrib(0, 255);

        for (size_t i = 0; i < size; ++i) {
            data[i] = static_cast<uint8_t>(distrib(gen));
        }

        return data;
    }
    
    // 计算对齐后的大小
    uint32_t AlignSize(uint32_t size, uint32_t alignment)
    {
        return (size + alignment - 1) & ~(alignment - 1);
    }
};

// 测试初始化
TEST_F(MemoryPoolTest, Initialization)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    
    // 测试初始化前的状态
    EXPECT_EQ(pool.GetTotalSize(), 0);
    EXPECT_EQ(pool.GetAvailableSize(), 0);
    EXPECT_EQ(pool.GetUsedSize(), 0);
    
    // 测试正常初始化
    EXPECT_TRUE(pool.Init(1024));
    EXPECT_EQ(pool.GetTotalSize(), 1024);
    EXPECT_EQ(pool.GetAvailableSize(), 1024);
    EXPECT_EQ(pool.GetUsedSize(), 0);
    
    // 测试重新初始化
    EXPECT_TRUE(pool.Init(2048, 16));
    EXPECT_EQ(pool.GetTotalSize(), 2048);
    EXPECT_EQ(pool.GetAvailableSize(), 2048);
    EXPECT_EQ(pool.GetUsedSize(), 0);
    
    // 测试初始化参数错误
    EXPECT_FALSE(pool.Init(0));
}

// 测试基本的内存分配和释放
TEST_F(MemoryPoolTest, BasicAllocateAndFree)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    // 分配内存
    void* ptr1 = pool.Allocate(100);
    ASSERT_NE(ptr1, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), 104); // 对齐到8字节
    EXPECT_EQ(pool.GetAvailableSize(), 1024 - 104);
    
    // 写入数据
    std::vector<uint8_t> testData = GenerateRandomData(100);
    memcpy(ptr1, testData.data(), 100);
    
    // 分配第二块内存
    void* ptr2 = pool.Allocate(200);
    ASSERT_NE(ptr2, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), 104 + 200); // 第一块对齐后100 + 第二块200
    EXPECT_EQ(pool.GetAvailableSize(), 1024 - 104 - 200);
    
    // 验证第一块数据未被破坏
    EXPECT_EQ(memcmp(ptr1, testData.data(), 100), 0);
    
    // 按顺序释放内存
    EXPECT_TRUE(pool.Free(ptr1));
    EXPECT_EQ(pool.GetUsedSize(), 200);
    
    EXPECT_TRUE(pool.Free(ptr2));
    EXPECT_EQ(pool.GetUsedSize(), 0);
    EXPECT_EQ(pool.GetAvailableSize(), 1024);
}

// 测试内存对齐
TEST_F(MemoryPoolTest, MemoryAlignment)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 16)); // 16字节对齐
    
    // 分配不同大小的内存，检查对齐
    void* ptr1 = pool.Allocate(10);
    ASSERT_NE(ptr1, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), 16); // 应该对齐到16字节
    
    void* ptr2 = pool.Allocate(16);
    ASSERT_NE(ptr2, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), 32); // 10对齐到16 + 16
    
    void* ptr3 = pool.Allocate(17);
    ASSERT_NE(ptr3, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), 64); // 10对齐到16 + 16 + 17对齐到32
    
    // 检查指针对齐
    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr1) % 16, 0);
    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr2) % 16, 0);
    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr3) % 16, 0);
    
    // 按顺序释放内存
    EXPECT_TRUE(pool.Free(ptr1));
    EXPECT_TRUE(pool.Free(ptr2));
    EXPECT_TRUE(pool.Free(ptr3));
}

// 测试顺序释放
TEST_F(MemoryPoolTest, OrderedFree)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    void* ptr1 = pool.Allocate(100);
    void* ptr2 = pool.Allocate(200);
    void* ptr3 = pool.Allocate(300);
    
    ASSERT_NE(ptr1, nullptr);
    ASSERT_NE(ptr2, nullptr);
    ASSERT_NE(ptr3, nullptr);
    
    // 按顺序释放，应该都成功
    EXPECT_TRUE(pool.Free(ptr1));
    EXPECT_TRUE(pool.Free(ptr2));
    EXPECT_TRUE(pool.Free(ptr3));
    
    // 所有内存都应该被释放
    EXPECT_EQ(pool.GetUsedSize(), 0);
    EXPECT_EQ(pool.GetAvailableSize(), 1024);
}

// 测试非顺序释放（现在应该成功）
TEST_F(MemoryPoolTest, NonOrderedFree)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    void* ptr1 = pool.Allocate(100);
    void* ptr2 = pool.Allocate(200);
    void* ptr3 = pool.Allocate(300);
    
    ASSERT_NE(ptr1, nullptr);
    ASSERT_NE(ptr2, nullptr);
    ASSERT_NE(ptr3, nullptr);
    
    // 尝试非顺序释放，现在应该成功
    EXPECT_TRUE(pool.Free(ptr2)); // 非顺序释放
    
    // 释放其他块
    EXPECT_TRUE(pool.Free(ptr1));
    EXPECT_TRUE(pool.Free(ptr3));
    
    // 所有内存都应该被释放
    EXPECT_EQ(pool.GetUsedSize(), 0);
    EXPECT_EQ(pool.GetAvailableSize(), 1024);
}

// 测试内存池容量限制和循环利用
TEST_F(MemoryPoolTest, CapacityLimitsAndReuse)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    // 分配接近容量的内存
    void* ptr1 = pool.Allocate(500);
    ASSERT_NE(ptr1, nullptr);
    
    void* ptr2 = pool.Allocate(500);
    ASSERT_NE(ptr2, nullptr);
    
    // 尝试分配超出剩余容量的内存
    void* ptr3 = pool.Allocate(100);
    EXPECT_EQ(ptr3, nullptr);
    
    // 释放第一块后再尝试，应该可以从头开始分配
    EXPECT_TRUE(pool.Free(ptr1));
    ptr3 = pool.Allocate(100);
    ASSERT_NE(ptr3, nullptr);
    
    // 释放所有内存
    EXPECT_TRUE(pool.Free(ptr2));
    EXPECT_TRUE(pool.Free(ptr3));
    
    // 重置后应该可以再次分配
    pool.Reset();
    ptr1 = pool.Allocate(1000);
    ASSERT_NE(ptr1, nullptr);
    EXPECT_TRUE(pool.Free(ptr1));
}

// 测试重置功能
TEST_F(MemoryPoolTest, Reset)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    const uint32_t blockAlign = 8;
    ASSERT_TRUE(pool.Init(1024, blockAlign));
    
    // 分配一些内存
    const uint32_t size1 = 300;
    const uint32_t size2 = 300;
    void* ptr1 = pool.Allocate(size1);
    void* ptr2 = pool.Allocate(size2);
    
    ASSERT_NE(ptr1, nullptr);
    ASSERT_NE(ptr2, nullptr);
    
    // 计算对齐后的总大小
    uint32_t alignedSize1 = AlignSize(size1, blockAlign);
    uint32_t alignedSize2 = AlignSize(size2, blockAlign);
    uint32_t totalAlignedSize = alignedSize1 + alignedSize2;
    
    // 验证使用的内存大小
    EXPECT_EQ(pool.GetUsedSize(), totalAlignedSize);
    
    // 重置内存池
    pool.Reset();
    
    // 验证重置后的状态
    EXPECT_EQ(pool.GetUsedSize(), 0);
    EXPECT_EQ(pool.GetAvailableSize(), 1024);
    
    // 重置后应该可以再次分配
    ptr1 = pool.Allocate(500);
    ASSERT_NE(ptr1, nullptr);
    EXPECT_EQ(pool.GetUsedSize(), AlignSize(500, blockAlign)); // 对齐到8字节
}

// 测试边界情况
TEST_F(MemoryPoolTest, EdgeCases)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    // 测试分配0字节
    void* ptr = pool.Allocate(0);
    EXPECT_EQ(ptr, nullptr);
    
    // 测试分配超过内存池大小的内存
    ptr = pool.Allocate(2000);
    EXPECT_EQ(ptr, nullptr);
    
    // 测试释放nullptr
    EXPECT_FALSE(pool.Free(nullptr));
    
    // 测试释放内存池外的指针
    uint8_t externalBuffer[100];
    EXPECT_FALSE(pool.Free(externalBuffer));
}

// 测试循环分配
TEST_F(MemoryPoolTest, CircularAllocation)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    // 分配一些内存，然后释放
    void* ptr1 = pool.Allocate(400);
    void* ptr2 = pool.Allocate(400);
    
    ASSERT_NE(ptr1, nullptr);
    ASSERT_NE(ptr2, nullptr);
    
    // 释放第一块
    EXPECT_TRUE(pool.Free(ptr1));
    
    // 尝试分配一个大块，这应该成功，因为第一块已释放，可以从头开始分配
    void* ptr3 = pool.Allocate(400);
    ASSERT_NE(ptr3, nullptr);
    
    // 释放所有内存
    EXPECT_TRUE(pool.Free(ptr2));
    EXPECT_TRUE(pool.Free(ptr3));
}

// 测试内存池的分配行为
TEST_F(MemoryPoolTest, Wraparound)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024, 8));
    
    // 分配几乎所有内存
    void* ptr1 = pool.Allocate(900);
    ASSERT_NE(ptr1, nullptr);
    
    // 再分配一个小块，应该成功
    void* ptr2 = pool.Allocate(100);
    ASSERT_NE(ptr2, nullptr);
    
    // 再分配一个块，这应该失败，因为没有足够的空间
    void* ptr3 = pool.Allocate(100);
    EXPECT_EQ(ptr3, nullptr);
    
    // 释放内存
    EXPECT_TRUE(pool.Free(ptr1));
    EXPECT_TRUE(pool.Free(ptr2));
    
    // 重置后应该可以再次分配
    pool.Reset();
    ptr1 = pool.Allocate(1000);
    ASSERT_NE(ptr1, nullptr);
    EXPECT_TRUE(pool.Free(ptr1));
}

// 测试多线程分配和释放
TEST_F(MemoryPoolTest, MultithreadedAllocateAndFree)
{
    MemoryPool& pool = MemoryPool::GetInstance();
    ASSERT_TRUE(pool.Init(1024 * 1024, 8)); // 使用更大的内存池
    
    const int numThreads = 4;
    const int allocationsPerThread = 100;
    const int allocationSize = 1000;
    
    std::vector<std::vector<void*>> threadPtrs(numThreads);
    std::vector<std::thread> threads;
    
    // 创建多个线程进行分配
    for (int i = 0; i < numThreads; ++i) {
        threads.emplace_back([i, allocationsPerThread, allocationSize, &pool, &threadPtrs]() {
            for (int j = 0; j < allocationsPerThread; ++j) {
                void* ptr = pool.Allocate(allocationSize);
                if (ptr != nullptr) {
                    threadPtrs[i].push_back(ptr);
                }
            }
        });
    }
    
    // 等待所有线程完成
    for (auto& thread : threads) {
        thread.join();
    }
    
    // 验证分配结果
    size_t totalAllocations = 0;
    for (const auto& ptrs : threadPtrs) {
        totalAllocations += ptrs.size();
    }
    
    // 注意：由于多线程环境，我们不能保证按顺序释放
    // 所以这里我们只是重置内存池
    pool.Reset();
    
    // 验证所有内存都已释放
    EXPECT_EQ(pool.GetUsedSize(), 0);
}

// 测试多实例功能
TEST_F(MemoryPoolTest, MultipleInstances)
{
    // 创建三个不同名称的内存池实例
    MemoryPool& pool1 = MemoryPool::GetInstance("pool1");
    MemoryPool& pool2 = MemoryPool::GetInstance("pool2");
    MemoryPool& pool3 = MemoryPool::GetInstance("pool3");
    
    // 初始化每个内存池，使用不同的大小
    ASSERT_TRUE(pool1.Init(1024));
    ASSERT_TRUE(pool2.Init(2048));
    ASSERT_TRUE(pool3.Init(4096));
    
    // 验证每个内存池的大小是否正确
    EXPECT_EQ(pool1.GetTotalSize(), 1024);
    EXPECT_EQ(pool2.GetTotalSize(), 2048);
    EXPECT_EQ(pool3.GetTotalSize(), 4096);
    
    // 在每个内存池中分配内存
    void* ptr1 = pool1.Allocate(256);
    void* ptr2 = pool2.Allocate(512);
    void* ptr3 = pool3.Allocate(1024);
    
    // 验证分配是否成功
    ASSERT_NE(ptr1, nullptr);
    ASSERT_NE(ptr2, nullptr);
    ASSERT_NE(ptr3, nullptr);
    
    // 验证每个内存池的可用大小是否正确
    EXPECT_EQ(pool1.GetAvailableSize(), 1024 - 256);
    EXPECT_EQ(pool2.GetAvailableSize(), 2048 - 512);
    EXPECT_EQ(pool3.GetAvailableSize(), 4096 - 1024);
    
    // 释放内存
    EXPECT_TRUE(pool1.Free(ptr1));
    EXPECT_TRUE(pool2.Free(ptr2));
    EXPECT_TRUE(pool3.Free(ptr3));
    
    // 验证释放后的可用大小
    EXPECT_EQ(pool1.GetAvailableSize(), 1024);
    EXPECT_EQ(pool2.GetAvailableSize(), 2048);
    EXPECT_EQ(pool3.GetAvailableSize(), 4096);
    
    // 测试获取相同名称的实例是否返回同一个实例
    MemoryPool& pool1Again = MemoryPool::GetInstance("pool1");
    EXPECT_EQ(&pool1, &pool1Again);
    
    // 测试释放单个实例
    EXPECT_TRUE(MemoryPool::ReleaseInstance("pool1"));
    
    // 再次获取应该是一个新的实例
    MemoryPool& pool1New = MemoryPool::GetInstance("pool1");
    EXPECT_NE(&pool1, &pool1New);
    
    // 释放所有实例
    MemoryPool::ReleaseAllInstances();
    
    // 再次获取应该是新的实例
    MemoryPool& pool2New = MemoryPool::GetInstance("pool2");
    EXPECT_NE(&pool2, &pool2New);
}